import logging
from datetime import datetime
from pinecone import Pinecone
from sklearn.datasets import load_digits
from tqdm import tqdm  # 用于进度条
import numpy as np
from collections import Counter
from sklearn.model_selection import train_test_split
from pinecone import Pinecone, ServerlessSpec

# 配置日志记录
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)

# Pinecone 配置
pinecone = Pinecone(api_key="d3e05f35-b35d-484e-946e-11f7af1074be")
index_name = "mnist-index"

# 获取现有索引列表
existing_indexes = pinecone.list_indexes()  # 获取现有索引列表


# 检查索引是否存在，如果存在就删除
'''
if any(index['name'] == index_name for index in existing_indexes):
    print(f"索引 '{index_name}' 已存在，正在删除...")
    pinecone.delete_index(index_name)
    print(f"索引 '{index_name}' 已成功删除。")
else:
    print(f"索引 '{index_name}' 不存在，将创建新索引。")

# 创建新索引
pinecone.create_index(
    name=index_name,
    dimension=64,
    metric='euclidean',
    spec=ServerlessSpec(
        cloud='aws',
        region='us-east-1'
    )
    
)
'''
index = pinecone.Index(index_name)
logging.info(f"已成功创建并连接到索引 '{index_name}'。")

# 加载 MNIST 数据集
digits = load_digits(n_class=10)
X = digits.data
y = digits.target

# 划分数据集，80% 用于训练，20% 用于测试
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 数据预处理
vectors = []
for i in tqdm(range(len(X)), desc="预处理数据"):
    vector_id = str(i)
    vector_values = X[i].tolist()
    metadata = {"label": int(y[i])}
    vectors.append((vector_id, vector_values, metadata))

# 定义批处理大小
batch_size = 1000

# 上传数据到 Pinecone
for i in range(0, len(vectors), batch_size):
    batch = vectors[i:i + batch_size]
    index.upsert(batch)
    logging.info(f"已上传批次数据，范围: {i}-{i + batch_size - 1}")

# 日志记录上传完成
logging.info(f"成功创建索引，并上传了{len(vectors)}条数据")

# 使用 20% 的数据进行测试
test_size = int(0.2 * len(vectors))
test_vectors = vectors[-test_size:]

# 测试 k=11 时的准确率
correct_predictions = 0
k = 11

for vector_id, vector_values, metadata in tqdm(test_vectors, desc="测试k=11准确率"):
    search_results = index.query(vector=vector_values, top_k=k, include_metadata=True)
    labels = [result.metadata['label'] for result in search_results['matches']]
    label_count = Counter(labels)
    predicted_label = label_count.most_common(1)[0][0]
    if predicted_label == metadata['label']:
        correct_predictions += 1

accuracy = correct_predictions / len(test_vectors)
logging.info(f"当k={k}时，使用Pinecone的准确率为: {accuracy:.2f}")


# 示例中创建手写数字图像的代码被省略，因为其与测试准确率无关